import fiftyone as fo
import fiftyone.brain as fob
from fiftyone import ViewField as F
import os
from pathlib import Path


def find_similar_images(
    dataset,
    query_identifier,
    top_k=50,
    similarity_field="similarity",  # 旧版本中使用的字段名
):
    # 检查brain索引并计算相似度
    brain_keys = dataset.list_brain_runs()
    if similarity_field not in brain_keys:
        print("未检测到相似度数据，开始计算...")
        fob.compute_similarity(dataset, brain_key=similarity_field)
    else:
        # similarity_index = dataset.get_brain_run("similarity")
        print("相似度数据已存在，跳过计算")

    # 规范化查询路径
    if os.path.exists(query_identifier):
        query_path = Path(query_identifier).resolve()
        query_identifier = str(query_path)

    # 确定查询样本ID
    if os.path.exists(query_identifier):
        matches = dataset.match(F("filepath") == query_identifier)
        if len(matches) == 0:
            matches = dataset.match(
                F("filepath").ends_with(os.path.basename(query_identifier))
            )
            if len(matches) == 0:
                raise ValueError(f"未在数据集中找到图像: {query_identifier}")
            elif len(matches) > 1:
                print(f"警告: 找到多个同名文件，将使用第一个匹配结果")
        sample_id = matches.first().id
    else:
        sample_id = query_identifier
        if not dataset.exists(sample_id):
            raise ValueError(f"样本ID不存在: {sample_id}")

    # 查找相似图像（旧版本方法）
    print(f"正在查找与样本 {sample_id} 最相似的 {top_k} 张图像...")
    view = dataset.sort_by_similarity(
        sample_id,
        brain_key=similarity_field,  # 旧版本使用brain_key参数
        k=top_k + 1,
    )

    similar_samples = [s for s in view if s.id != sample_id][:top_k]
    return sample_id, similar_samples


if __name__ == "__main__":
    DATASET_NAME = "carriage_ocr_dataset"
    QUERY_IMAGE = "chars/24_20250731112436917_e6063c33908e4dc4a3f81cd6a7384439.jpeg"
    TOP_K = 10

    try:
        print(f"正在加载数据集: {DATASET_NAME}")
        dataset = fo.load_dataset(DATASET_NAME)
        print(f"数据集加载成功，包含 {len(dataset)} 个样本")

        print(dataset)

        query_id, similar_images = find_similar_images(
            dataset,
            QUERY_IMAGE,
            top_k=TOP_K,
        )

        print(f"\n与查询图像最相似的 {TOP_K} 个图像:")
        for i, sample in enumerate(similar_images, 1):
            print(f"{i}. {sample.filepath}")

        # 标记样本
        # dataset.match_tags("query").clear_tags("query")
        # dataset.match_tags("similar").clear_tags("similar")

        query_sample = dataset[query_id]
        query_sample.tags.append("query")
        query_sample.save()

        for sample in similar_images:
            sample.tags.append("similar")
            sample.save()

        results_view = dataset.select([query_id] + [s.id for s in similar_images])
        print("启动FiftyOne可视化界面...")
        session = fo.launch_app(results_view)
        session.filters = {"tags": {"$in": ["query"]}}
        session.wait()

    except fo.DatasetNotFoundError:
        print(f"错误: 数据集 {DATASET_NAME} 不存在，请先创建并导入图像")
    except ValueError as e:
        print(f"错误: {e}")
    except Exception as e:
        print(f"程序运行出错: {str(e)}")
